"""
Copyright (c) 2022 Beijing Jiaotong University
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Chunyu Li
Created: 2022/9/6
Supported by: National Key Research and Development Program of China
"""

import numpy as np
from ...optical import fine_synchronize, ber_count
import commpy as cp

look_up_table_decoded_8QAM = [[1, 1, 0], [0, 1, 0], [1, 0, 0], [0, 0, 0], [0, 1, 1], [1, 0, 1], [1, 1, 1], [0, 0, 1]]

mod_table_8QAM = [-1-1j, -1+1j, 1-1j, 1+1j, -1-np.sqrt(3), 1+np.sqrt(3), (-1-np.sqrt(3))*1j, (1+np.sqrt(3))*1j]

look_up_table_decoded_32QAM = [[1, 0, 1, 1, 0], [1, 0, 1, 1, 1], [1, 0, 1, 0, 1], [1, 0, 1, 0, 0],
     [1, 0, 0, 1, 0], [1, 1, 1, 1, 0], [1, 1, 1, 1, 1], [1, 1, 1, 0, 1], [1, 1, 1, 0, 0], [1, 0, 0, 0, 0],
     [1, 0, 0, 1, 1], [1, 1, 0, 1, 0], [1, 1, 0, 1, 1], [1, 1, 0, 0, 1], [1, 1, 0, 0, 0], [1, 0, 0, 0, 1],
     [0, 0, 0, 1, 1], [0, 1, 0, 1, 0], [0, 1, 0, 1, 1], [0, 1, 0, 0, 1], [0, 1, 0, 0, 0], [0, 0, 0, 0, 1],
     [0, 0, 0, 1, 0], [0, 1, 1, 1, 0], [0, 1, 1, 1, 1], [0, 1, 1, 0, 1], [0, 1, 1, 0, 0], [0, 0, 0, 0, 0],
     [0, 0, 1, 1, 0], [0, 0, 1, 1, 1], [0, 0, 1, 0, 1], [0, 0, 1, 0, 0]]

mod_table_32QAM = \
    [5 + 3j, 5 + 1j, 5 - 1j, 5 - 3j, 3 + 5j, 3 + 3j, 3 + 1j, 3 - 1j, 3 - 3j, 3 - 5j, 1 + 5j,
     1 + 3j, 1 + 1j, 1 - 1j, 1 - 3j, 1 - 5j, -1 + 5j, -1 + 3j, -1 + 1j, -1 - 1j, -1 - 3j, -1 - 5j,
     -3 + 5j, -3 + 3j, -3 + 1j, -3 - 1j, -3 - 3j, -3 - 5j, -5 + 3j, -5 + 1j, -5 - 1j, -5 - 3j]


#TODO:该函数将来会作为一个独立组件
def deModulation(signals, bits_per_symbol):
    if bits_per_symbol == 2 or bits_per_symbol == 4 or bits_per_symbol == 6:
        qam64 = cp.QAMModem(2 ** bits_per_symbol)
        x = qam64.demodulate(signals[0].reshape(-1), 'hard')
        y = qam64.demodulate(signals[1].reshape(-1), 'hard')
    elif bits_per_symbol == 3:
        X, Y = signals[0].reshape(-1), signals[1].reshape(-1)
        x, y = [], []
        for i in range(len(X)):
            x.append(look_up_table_decoded_8QAM[np.argmin(np.abs((X[i] - np.array(mod_table_8QAM))**2))])
        for j in range(len(Y)):
            y.append(look_up_table_decoded_8QAM[np.argmin(np.abs((Y[j] - np.array(mod_table_8QAM))**2))])
    elif bits_per_symbol == 5:
        X, Y = signals[0].reshape(-1), signals[1].reshape(-1)
        x, y = [], []
        for i in range(len(X)):
            x.append(look_up_table_decoded_32QAM[np.argmin(np.abs((X[i] - np.array(mod_table_32QAM))**2))])
        for j in range(len(Y)):
            y.append(look_up_table_decoded_32QAM[np.argmin(np.abs((Y[j] - np.array(mod_table_32QAM))**2))])

    x, y = np.expand_dims(np.array(x).reshape(-1), axis=1), np.expand_dims(np.array(y).reshape(-1), axis=1)
    return [x, y]


def Correct_phase(input, traning_input, Bits_per_symbol):
    """纠正相位模糊
            Args:
                input[0]: 输入X偏振信号, 尺寸为(n, 1)
                input[1]: 输入Y偏振信号, 尺寸为(n, 1)
                traning_input[0]: 输入X偏振训练样本，尺寸为(n, 1)
                traning_input[1]: 输入Y偏振训练样本，尺寸为(n, 1)
                Bits_per_symbol：比特数

            Returns:
                output[0]: 输出X偏振信号
                output[1]: 输出Y偏振信号
                traning_output[0]: 输出X偏振训练样本
                traning_output[1]: 输出Y偏振训练样本
            """
    Equalization_matrix_x = input[0].reshape((-1, 1))
    Equalization_matrix_y = input[1].reshape((-1, 1))
    Traning_Sample_X = traning_input[0].reshape((-1, 1))
    Traning_Sample_Y = traning_input[1].reshape((-1, 1))

    Traning_Sample_X = Traning_Sample_X * np.exp(-1j * np.angle(Traning_Sample_X * np.conj(Equalization_matrix_x)))
    Traning_Sample_Y = Traning_Sample_Y * np.exp(-1j * np.angle(Traning_Sample_Y * np.conj(Equalization_matrix_y)))

    output =deModulation([Equalization_matrix_x,Equalization_matrix_y], Bits_per_symbol)
    traning_output =deModulation([Traning_Sample_X,Traning_Sample_Y], Bits_per_symbol)

    return [output, traning_output]


def bits_error_count(signals, bits_per_symbol):
    """计算误码率"""

    signal_x = signals[0].reshape((-1, 1))
    signal_y = signals[1].reshape((-1, 1))
        
    prev_symbols_x = signals[2].reshape((-1, 1))
    prev_symbols_y = signals[3].reshape((-1, 1))

    """ 再进行一个帧同步，因为经过均衡器会存在符号的一些舍弃，因此在计算误码率（BER）之前需要再一次帧同步 """

    # 对发射端信号跟均衡后信号进行同步
    start_index_x_1 = fine_synchronize(prev_symbols_x[:, 0].T, signal_x[0:10000, 0].reshape((1, -1)))

    # 对发射端信号利用自带函数进行移动
    prev_symbols_x = np.roll(prev_symbols_x, -start_index_x_1)
    prev_symbols_y = np.roll(prev_symbols_y, -start_index_x_1)

    # 变为16的倍数
    signal_x = signal_x[:-1]
    signal_y = signal_y[:-1]

    # 使得均衡后信号与发射端信号一样长度
    prev_symbols_x = prev_symbols_x[0 : len(signal_x), :]
    prev_symbols_y = prev_symbols_y[0 : len(signal_y), :]

    signal_x, prev_symbols_x = Correct_phase([signal_x, signal_y], [prev_symbols_x, prev_symbols_y], bits_per_symbol)

    """ BER COUNT  对信号进行误码率计算，将接收信号与发射信号转化为格雷编码，比较各个码元的正确率 """

    # ber1, q_factor1 = ber_estimate(prev_symbols_x[0], signal_x[0], bits_per_symbol)
    # ber, q_factor = ber_estimate(prev_symbols_x[1], signal_x[1], bits_per_symbol)

    our_ber_0 = ber_count(signal_x[0], prev_symbols_x[0])
    our_ber_1 = ber_count(signal_x[1], prev_symbols_x[1])

    # return (ber+ber1)/2, (q_factor+q_factor1)/2,our_ber
    # return 0.25 * (our_ber_0[0] + our_ber_0[1] + our_ber_1[0] + our_ber_1[1])

    return 0.5 * (our_ber_0[0] + our_ber_1[0]), 0.5 * (our_ber_0[1] + our_ber_1[1])